#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import argparse
import logging
from scapy.all import sniff, IP
from datetime import datetime
import json
import os

from .rag_detector import RAGVPNDetector
from .training_data_generator import TrainingDataGenerator

def setup_logging():
    """设置日志"""
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )

def save_detection_results(results: list, output_file: str):
    """保存检测结果"""
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(results, f, indent=2, ensure_ascii=False)

def main():
    """主程序"""
    parser = argparse.ArgumentParser(description='VPN流量检测系统')
    parser.add_argument('--interface', type=str, required=True, help='网络接口名称')
    parser.add_argument('--mode', type=str, choices=['detect', 'train'], required=True, help='运行模式：detect(检测) 或 train(训练)')
    parser.add_argument('--duration', type=int, default=300, help='捕获时长(秒)')
    parser.add_argument('--output', type=str, default='detection_results.json', help='输出文件路径')
    
    args = parser.parse_args()
    setup_logging()
    logger = logging.getLogger(__name__)
    
    if args.mode == 'train':
        # 训练模式
        logger.info("开始训练模式")
        generator = TrainingDataGenerator()
        generator.generate_all_protocols_data(args.interface, args.duration)
        logger.info("训练完成")
        
    else:
        # 检测模式
        logger.info("开始检测模式")
        detector = RAGVPNDetector()
        results = []
        
        def packet_callback(packet):
            if packet.haslayer(IP):
                result = detector.detect(packet)
                if result.protocol:
                    results.append({
                        'timestamp': datetime.now().isoformat(),
                        'protocol': result.protocol.value,
                        'confidence': result.confidence,
                        'features': result.features,
                        'evidence': result.evidence
                    })
                    
        logger.info(f"开始在接口 {args.interface} 上捕获流量")
        sniff(iface=args.interface, prn=packet_callback, timeout=args.duration)
        
        # 保存结果
        save_detection_results(results, args.output)
        logger.info(f"检测完成，结果已保存到 {args.output}")

if __name__ == '__main__':
    main() 